package com.shujia.sql

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo05Burks {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.replace("$", ""))
      .master("local")
      .config("spark.sql.shuffle.partitions", "2")
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._

    val burkDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", ",")
      .schema("burk String,year String,tsl01 String,tsl02 String,tsl03 String,tsl04 String,tsl05 String,tsl06 String,tsl07 String,tsl08 String,tsl09 String,tsl10 String,tsl11 String,tsl12 String")
      .load("spark/data/burks.txt")

    burkDF.show()
    burkDF.createOrReplaceTempView("burk_tb")
    // SQL的方式
    spark.sql(
      """
        |select  t1.burk
        |        ,t1.year
        |        ,t1.month
        |        ,t1.tsl
        |        ,sum(t1.tsl) over (partition by t1.burk,t1.year order by t1.month) as sum_tsl
        |from (
        |    select  burk
        |            ,year
        |            ,(pos + 1) as month
        |            ,tsl
        |    from burk_tb
        |    lateral view posexplode(array(tsl01,tsl02,tsl03,tsl04,tsl05,tsl06,tsl07,tsl08,tsl09,tsl10,tsl11,tsl12)) v1 as pos,tsl
        |) t1
        |""".stripMargin)
    //      .show()

    // DSL的方式
    val yearMonthTslDF: DataFrame = burkDF
      .select($"burk", $"year", posexplode(array($"tsl01", $"tsl02", $"tsl03", $"tsl04", $"tsl05", $"tsl06", $"tsl07", $"tsl08", $"tsl09", $"tsl10", $"tsl11", $"tsl12")))
      .select($"burk", $"year", $"pos" + 1 as "month", $"col" as "tsl")

    // 多次使用的DF也可以进行缓存
    yearMonthTslDF.cache()

    yearMonthTslDF
      //      .select($"burk", $"year", $"month", $"tsl", sum($"tsl")
      //        .over(Window.partitionBy($"burk", $"year").orderBy($"month")) as "sum_tsl")
      .withColumn("sum_tsl", sum($"tsl")
        .over(Window.partitionBy($"burk", $"year").orderBy($"month")))
    //      .show()

    yearMonthTslDF
      .withColumn("last_tsl", lag($"tsl", 1) over Window.partitionBy($"burk", $"month").orderBy($"year"))
      .withColumn("tb", coalesce(round($"tsl" / $"last_tsl" - 1, 40), expr("0")))
      .show()

  }

}
